feat(pt/dpa4): add nvalchemi-toolkit warpper for DPA4#5568
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (11)
💤 Files with no reviewable changes (1)
✅ Files skipped from review due to trivial changes (3)
🚧 Files skipped from review as they are similar to previous changes (6)
📝 WalkthroughWalkthroughAdds ChangesDPA4Wrapper nvalchemi adapter
Sequence Diagram(s)sequenceDiagram
participant User as User script
participant DPA4Wrapper
participant NeighborListHook
participant SeZMModel
User->>DPA4Wrapper: from_checkpoint(path, device, ...)
DPA4Wrapper-->>User: wrapper instance
User->>NeighborListHook: register before-compute hook
User->>DPA4Wrapper: forward(batch)
DPA4Wrapper->>DPA4Wrapper: _atype() and _edge_schema()
DPA4Wrapper->>DPA4Wrapper: adapt_input()
DPA4Wrapper->>SeZMModel: forward_lower(...) / AOT callable
SeZMModel-->>DPA4Wrapper: atom_energy, extended_force, extended_virial
DPA4Wrapper->>DPA4Wrapper: adapt_output()
DPA4Wrapper-->>User: energy, forces, stress
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/nvalchemi/wrapper.py`:
- Line 485: Before accessing model_params["model_dict"][head] at line 485,
validate that the head parameter exists as a key in the model_dict dictionary.
If the head is not found, raise a ValueError (instead of letting a raw KeyError
propagate) that includes a descriptive message with the list of available heads
to provide better debugging information and a cleaner API contract.
- Around line 274-283: The conditional check on line 274 that evaluates `if
neighbor_list_shifts is not None and cell is not None` allows incomplete
periodic metadata to pass through silently, causing incorrect physics when one
field is provided without the other. Add validation logic before this
conditional to check that both neighbor_list_shifts and cell are either both
present or both absent, and raise an informative error if they are inconsistent
(one is None while the other is not). This ensures that periodic boundary
condition calculations fail fast with a clear error message rather than silently
using incorrect zero shifts.
- Around line 366-367: The stress computation at the line where virial is
divided by volume does not validate whether the volume is valid before
performing the division, which can result in inf or nan values for singular or
degenerate cells. Add a validation check before computing the stress output that
ensures the volume tensor contains only valid (non-zero, non-NaN, non-inf)
values, and raise a clear and informative error if any invalid volumes are
detected. This check should occur after computing the volume using
torch.det(cell.to(virial.dtype)).abs().view(n_graph, 1, 1) but before the
division in the stress calculation.
In `@examples/water/dpa4/nvalchemi/relax.py`:
- Around line 95-96: The argument parser for --log-every lacks validation and
the optimization loop does not properly handle remaining steps. Add validation
to the argument parser around the --log-every definition to ensure the value is
greater than 0. Then examine the stepping logic around lines 141-151 where the
loop increments by log_every chunks, and modify it to ensure all remaining steps
up to --max-steps are executed even when --max-steps is not perfectly divisible
by log_every, preventing under-run of the optimization budget.
In `@examples/water/dpa4/nvalchemi/run_nve.py`:
- Around line 114-115: Add validation for the --log-every argument to ensure it
is greater than zero, preventing runtime crashes from zero or negative values.
Additionally, fix the stepping logic in the simulation loop to handle cases
where the total number of steps is not evenly divisible by log_every. Instead of
only advancing in fixed log_every chunks, explicitly run any remainder steps at
the end of the loop to ensure all requested steps are executed (e.g., if 205
steps are requested with log_every 20, ensure 5 additional steps are run after
the main loop completes). Update the argument parser to include validation for
--log-every and modify the stepping logic around the main simulation iteration
to calculate and execute remaining steps.
In `@examples/water/dpa4/nvalchemi/run_nvt.py`:
- Around line 104-105: Add validation for the log_every argument to ensure it is
a positive integer (greater than zero) in the argument parser definition.
Additionally, fix the NVT integration loop around line 168-174 to properly
handle cases where the total number of steps is not evenly divisible by
log_every. Ensure the remainder steps are executed after the main chunked loop
completes so that the simulation runs for the full requested number of steps
rather than truncating early.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8c562042-bbdf-4ddd-806e-688cf8ca0c38
📒 Files selected for processing (11)
deepmd/pt/nvalchemi/__init__.pydeepmd/pt/nvalchemi/wrapper.pydoc/inference/index.rstdoc/inference/nvalchemi.mdexamples/water/dpa4/nvalchemi/README.mdexamples/water/dpa4/nvalchemi/relax.pyexamples/water/dpa4/nvalchemi/run_nve.pyexamples/water/dpa4/nvalchemi/run_nvt.pyexamples/water/dpa4/nvalchemi/single_point.pypyproject.tomlsource/tests/pt/model/test_sezm_nvalchemi.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5568 +/- ##
==========================================
- Coverage 82.27% 82.26% -0.02%
==========================================
Files 882 884 +2
Lines 100009 100230 +221
Branches 4045 4045
==========================================
+ Hits 82283 82452 +169
- Misses 16267 16317 +50
- Partials 1459 1461 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
OutisLi
force-pushed
the
pr/nvtoolkit-dpa4
branch
from
11ffa62 to
7ea521f
Compare
PR Summary by QodoAdd nvalchemi-toolkit adapter for DPA-4/SeZM PyTorch models Description
Diagram
High-Level Assessment
Files changed (11)
|
Code Review by Qodo
1. Sparse edges passed as nlist
|
![qodo-free-for-open-source-projects[bot]](https://avatars.githubusercontent.com/in/1068000?s=60&v=4){kind=small}
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Thanks for adding this integration. I think this needs changes before merge: I verified the current head (f62b88775cfb00ca616bea857463c13002c48320) and several main paths in the wrapper are still incompatible with the SeZM lower/AOTI interfaces, so the nvalchemi adapter would fail at runtime rather than just needing polish.
Blocking issues:
-
DPA4Wrapper.forward()passes sparse-edge tensors intoSeZMModel.forward_lower()with the wrong interface.
Indeepmd/pt/nvalchemi/dpa4wrapper.py, the wrapper buildsedge_indexas a COO tensor with shape(2, E)and passesedge_index,edge_vec,edge_scatter_index, andedge_maskpositionally toforward_lower. But the currentSeZMModel.forward_lower()signature expects DeePMD lower-interface inputs:(extended_coord, extended_atype, nlist, mapping, fparam, aparam, ...), wherenlistis a padded 3D neighbor list(nf, nloc, nsel). The model immediately routes throughformat_nlist()/build_edge_list_from_nlist(), so a(2, E)COO tensor will not work. Please either convert the nvalchemi COO list into the standard paddednlist+mapping, or add a supported sparse-edge SeZM entrypoint and call that explicitly. -
The
.pt2backend has the same callable-signature mismatch.
_run_pt2()calls the AOTI runner with 9 sparse-edge-style arguments, but the existing export path/test for DPA-4.pt2artifacts uses the 7-argumentforward_common_lowerinterface:(ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin). A standarddp --pt freezepackage will therefore not match this wrapper call. This should be fixed consistently with item 1, and ideally covered with a.pt2wrapper test. -
compute_embeddings()calls an unsupported API.
self.model.forward_common_lower(..., embedding_only=True)is not a valid call in this repository:forward_common_lower()has noembedding_onlykeyword, and it does not return adescriptorkey. Please implement embeddings via a supported descriptor path (for example, usingatomic_model.descriptor.forward_with_edges(...)with the prepared edge tensors), or add a dedicated SeZM method for descriptor-only output. -
The new optional extra is unconditional but
nvalchemi-toolkitis not available on all DeePMD-supported environments.
pyproject.tomladdsnvalchemi = ["nvalchemi-toolkit"], whilenvalchemi-toolkit==0.1.0currently declares Python>=3.11,<3.14; DeePMD-kit still supports Python 3.10. The existing PyTorch extra already guardsnvalchemi-toolkit-opswith environment markers for Python/Linux. Please add appropriate markers for this extra and document the supported Python/platform constraints indoc/third-party/nvalchemi.mdand the example README. -
The documented example install path misses an example dependency.
The example scripts import ASE (from ase.data import atomic_numbers as ASE_Z), butpip install deepmd-kit[nvalchemi]only pullsnvalchemi-toolkit. A user following the docs/README will hitModuleNotFoundError: No module named 'ase'before reaching the wrapper. Either remove the ASE dependency from these scripts by using an already-required element mapping, or include/document ASE (or an appropriate nvalchemi extra) in the prerequisites.
Smaller follow-up: the stress path divides by abs(det(cell)) without validating finite/non-zero volume. It would be safer to fail fast with a clear ValueError for singular/invalid cells, especially before using the output in barostats.
The new tests are a good direction, but because nvalchemi-toolkit is optional and absent from the current CI environment, the wrapper tests appear to be skipped and these runtime interface issues are not exercised. It would be worth adding at least one CI-covered path (or a lightweight mocked/interface test) that catches the SeZM lower-interface call shape.
Reviewed by OpenClaw 2026.6.8 (844f405)
model: custom-chat-jinzhezeng-group/gpt-5.5
|
Should be merged after PR #5571 is merged, this PR relies on that |
OutisLi
force-pushed
the
pr/nvtoolkit-dpa4
branch
from
f62b887 to
8f2a9f4
Compare
| volume = torch.det(cell.to(virial.dtype)).abs().view(n_graph, 1, 1) | ||
| output["stress"] = (virial.view(n_graph, 3, 3) / volume).to(out_dtype) |
There was a problem hiding this comment.
Following up on the sign: the current parity check (stress * volume == native_virial) is circular — it just re-asserts the wrapper's own virial/volume definition, so it cannot catch a wrong sign or the missing symmetrization. A non-circular UT would settle this for good. Two options:
A. Finite-difference of energy w.r.t. strain (independent ground truth). Stress is the energy derivative under cell strain, derivable from energies alone — fully independent of how the wrapper computes virial:
# periodic system, model run in float64
def E(eps): # eps: symmetric 3x3 strain
F = I + eps
return model_energy(coord @ F.T, cell @ F.T) # affine-deform atoms + cell
dE_deps[i, j] = (E(+d * E_ij) - E(-d * E_ij)) / (2 * d) # central diff, d ~ 1e-4
stress_ref = sign * dE_deps / det(cell)
torch.testing.assert_close(wrapper_stress, stress_ref)Energy is a sign-unambiguous scalar, so dE/deps is the physically-correct stress; a flipped sign (or the missing symmetrization) fails the assert. Choosing the sign that matches nvalchemi's documented convention — and writing it explicitly — is exactly what resolves the open question, instead of asserting it in a comment.
B. Compare against deepmd's own stress (simpler). deepmd's ASE calculator already produces a Cauchy stress with a separately-tested convention: stress = -0.5 * (v + v.T) / V (deepmd/calculator.py:160-161). Assert the wrapper's stress equals that for the same structure. This is non-circular because the reference is deepmd's trusted, independently-validated convention — not the wrapper's virial/volume formula — so a sign flip (or asymmetry) relative to deepmd fails immediately.
Both need a periodic system with non-zero stress (the water example works), float64, and can be gated @skipUnless(NVALCHEMI_AVAILABLE) like the rest — now actually runnable in CI since nvalchemi-toolkit was added to the test deps. (A) is the gold standard (catches a bug even if deepmd and the wrapper shared one); (B) is a few lines and still catches a sign flip vs deepmd.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
I reviewed the current head (b4a9fcc) of this PR. The earlier interface concerns I checked against the current tree look addressed:
DPA4Wrappernow drives SeZM through the explicit sparse-edge lower interface, andforward_common_lower(..., embedding_only=True)/ descriptor output are present in the currentSeZMModelimplementation.- The wrapper now fails fast for periodic batches missing
neighbor_list_shifts, validates multi-taskhead, and handles non-divisible--steps/--log-everychunks in the examples. nvalchemi-toolkitis included in both thetestandnvalchemiextras under the current platform/Python markers.- The stress sign convention has a non-circular strain finite-difference test in addition to parity checks.
I do not see additional blocking issues from this pass. CI is green on the current revision.
— OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
I do not see a remaining blocking issue.
Two small non-blocking suggestions:
- Please consider guarding stress computation against non-finite or near-zero cell volumes before dividing by
det(cell). - Since
.pt2support is user-facing, a minimal.pt2wrapper smoke test would be helpful to pin the AOTI callable signature.
LGTM otherwise.
5 participants
Summary by CodeRabbit
Release Notes
New Features
DPA4Wrapper/SeZMWrapperintegration enabling energy/force (and optional stress) inference via nvalchemi-toolkit, including support for frozen.pt2packages.Documentation
Bug Fixes / Improvements
Tests